#pragma once

#include "file_comm.h"
#include "xxhash.h"
#include "Socket.h"
#include <ios>
#include <fstream>
#include <filesystem>
#include <mutex>
#include <map>


class FileReceiver {
public:
    FileReceiver(const char *path, const char *host, int port) :
            path(path), server(host, port) {}

    void recvFile() {
        server.listen([&](Socket sock) {
            OperatorType op;
            sock.recv((char *) &op, sizeof(op));
            if (op == COLLECT) {
                recvCollect(sock);
            } else if (op == CHUNK) {
                recvChunk(sock);
            }
        });
    }

private:
    std::filesystem::path getTempDirectory() {
        std::filesystem::path file_path(path);
        std::filesystem::path directory = file_path.parent_path() / "temp";
        if (!std::filesystem::exists(directory)) {
            std::filesystem::create_directory(directory);
        }
        return directory;
    }

    void recvChunk(const Socket &sock) {
        Chunk chunk{};
        sock.recv((char *) &chunk.head, sizeof(chunk.head));
        chunk.data = std::make_unique<char[]>(chunk.head.size);
        sock.recv(chunk.data.get(), chunk.head.size);
        uint32_t hash = XXHash::xxh32(chunk.data.get(), chunk.head.size, 0);
        char res;
        if (hash == chunk.head.hash) {
            std::unique_lock<std::mutex> lock(mtx);
            std::filesystem::path directory = getTempDirectory();
            std::ofstream file(directory / std::to_string(chunk.head.no), std::ios::binary);
            file.write(chunk.data.get(), chunk.head.size);
            res = 1;
            LOG("接收文件块 %d / %d 成功！\n", chunk.head.no, collect.count);
            result[chunk.head.no] = true;
        } else {
            res = 0;
            LOG("接收文件块 %d / %d 失败！\n", chunk.head.no, collect.count);
        }
        sock.send(&res, 1);
    }

    void recvCollect(const Socket &sock) {
        sock.recv((char *) &collect, sizeof(collect));

        std::ofstream out(path, std::ios::binary);
        std::filesystem::path directory = getTempDirectory();
        std::error_code ec;
        for (int no = 1; no <= collect.count; no++) {
            while (true) {
                {
                    std::unique_lock<std::mutex> lock(mtx);
                    if (result[no]) {
                        break;
                    }
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
            }

            std::ifstream file(directory / std::to_string(no), std::ios::binary);
            char buf[CHUNK_SIZE];
            file.read(buf, CHUNK_SIZE);
            out.write(buf, file.gcount());
            file.close();
            std::filesystem::remove(directory / std::to_string(no), ec);
        }
        std::filesystem::remove(directory, ec);
        server.close();
    }

private:
    const char *path;
    Server server;
    Collect collect{};
    std::mutex mtx;
    std::map<int, bool> result;
};